"""I3D backbone"""

import mindspore.nn as nn
from mindspore.ops import operations as P
from mindvideo.common.utils.class_factory import ClassFactory, ModuleType
from mindvideo.models.builder import build_backbone


class Unit3D(nn.Cell):
    """
    Unit3D definition.

    Args:
        in_channels (int): The number of channels of input frame images.
        out_channels(int): The number of channels of output frame images.
        kernel_shape (tuple): The size of the kernel (default (1,1,1)).
        stride (tuple): An integer or tuple/list of a single integer, specifying the
            strides of the pooling operation.
        padding (int): An integer, Padding size.
        activation_fn (): Activation function applied to each layer.
        use_batch_norm (bool): Whether to use BatchNorm3d.
        has_bias (bool): Whether to use Bias.

    Returns:
        Tensor, output tensor.

    Examples:
        >>> Unit3D(in_channels=64, out_channels=64, kernel_shape=(1, 1, 1), padding=0)
    """

    def __init__(self, in_channels,
               out_channels,
               kernel_shape=(1, 1, 1),
               stride=(1, 1, 1),
               padding=0,
               activation_fn=nn.ReLU(),
               use_batch_norm=True,
               has_bias=False):
        super(Unit3D, self).__init__()

        self.out_channels = out_channels
        self.kernel_shape = kernel_shape
        self.stride = stride
        self.use_batch_norm = use_batch_norm
        self.activation_fn = activation_fn
        self.has_bias = has_bias
        self.padding = padding
        self.conv3d = nn.Conv3d(in_channels=in_channels,
                                out_channels=self.out_channels,
                                kernel_size=self.kernel_shape,
                                stride=self.stride,
                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function
                                has_bias=self.has_bias)

        if self.use_batch_norm:
            self.bn = nn.BatchNorm3d(self.out_channels, eps=0.001, momentum=0.01)

    def construct(self, x):

        x = self.conv3d(x)
        if self.use_batch_norm:
            x = self.bn(x)
        if self.activation_fn is not None:
            x = self.activation_fn(x)
        return x

@ClassFactory.register(ModuleType.BACKBONE)
class InceptionI3d(nn.Cell):
    """
    InceptionI3d architecture.

    Args:
        in_channels (int): The number of channels of input frame images(default 3).
    Returns:
        Tensor, output tensor.

    Examples:
        >>> InceptionI3d(**kwargs)
    """

    def __init__(self,
                 in_channels=3):

        super(InceptionI3d, self).__init__()

        self.conv3d_1a_7x7 = Unit3D(in_channels=in_channels, out_channels=64, kernel_shape=(7, 7, 7),
                                            stride=(2, 2, 2), padding=(3,3,3))
        self.maxpool3d_2a_3x3 = P.MaxPool3D(kernel_size=(1, 3, 3), strides=(1, 2, 2),  pad_mode="same")

        self.conv3d_2b_1x1 = Unit3D(in_channels=64, out_channels=64, kernel_shape=(1, 1, 1), padding=0)

        self.conv3d_2c_3x3 = Unit3D(in_channels=64, out_channels=192, kernel_shape=(3, 3, 3), padding=1)

        self.maxpool3d_3a_3x3 = P.MaxPool3D(kernel_size=(1, 3, 3), strides=(1, 2, 2),  pad_mode="same")

        self.mixed_3b = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 192,
                                        "out_channels": [64,96,128,16,32,32]})

        self.mixed_3c = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 256,
                                        "out_channels": [128,128,192,32,96,64]})

        self.maxpool3d_4a_3x3 = P.MaxPool3D(kernel_size=(3, 3, 3), strides=(2, 2, 2),  pad_mode="same")

        self.mixed_4b = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 128+192+96+64,
                                        "out_channels": [192,96,208,16,48,64]})

        self.mixed_4c = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 192+208+48+64,
                                        "out_channels": [160,112,224,24,64,64]})

        self.mixed_4d = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 160+224+64+64,
                                        "out_channels": [128,128,256,24,64,64]})

        self.mixed_4e = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 128+256+64+64,
                                        "out_channels": [112,144,288,32,64,64]})

        self.mixed_4f = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 112+288+64+64,
                                        "out_channels": [256,160,320,32,128,128]})

        self.maxpool3d_5a_2x2 = P.MaxPool3D(kernel_size=(2, 2, 2), strides=(2, 2, 2),  pad_mode="same")

        self.mixed_5b = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 256+320+128+128,
                                        "out_channels": [256,160,320,32,128,128]})

        self.mixed_5c = build_backbone({"type": "Inception3dModule",
                                        "in_channels": 256+320+128+128,
                                        "out_channels": [384,192,384,48,128,128]})

        # self.avg_pool = AvgPool3D(kernel_size=(2, 7, 7),strides=(1, 1, 1))
        self.mean_op = P.ReduceMean(keep_dims=True)
        self.concat_op = P.Concat(axis=2)
        self.stridedslice_op = P.StridedSlice()

    def avgpool3d(self, x, kernel_size, strides):
        """
        Function avgpool3d.

        Args:
            x (): .
            kernel_size (tuple): The size of the kernel.
            strides (tuple): An integer or tuple/list of a single integer, specifying the
            strides of the pooling operation.
        Returns:
            Tensor, output tensor.

        Examples:
            >>> avgpool3d(x, kernel_size=(2,7,7), strides=(1,1,1))
        """

        # MindSpore1.5 has no AvgPool3D grad operation for gpu,
        # Please replace this function with real AvgPool3D in the future version of mindspore.
        n, c, in_d, in_height, in_width = x.shape

        out_d = (in_d - kernel_size[0]) // strides[0] + 1
        out_height = (in_height - kernel_size[1]) // strides[1] + 1
        out_width = (in_width - kernel_size[2]) // strides[2] + 1
        out = []
        for i in range(out_d):
            for j in range(out_height):
                for k in range(out_width):
                    start_i = i * strides[0]
                    start_j = j * strides[1]
                    start_k = k * strides[2]
                    end_i = start_i + kernel_size[0]
                    end_j = start_j + kernel_size[1]
                    end_k = start_k + kernel_size[2]
                    out.append(self.mean_op(
                               self.stridedslice_op(x, (0,0,start_i,start_j,start_k),
                               (n,c,end_i,end_j,end_k), (1,1,1,1,1)),(2,3,4)))
        out = self.concat_op(out)
        return out

    def construct(self, x):
        """Average pooling 3D construct."""
        x = self.conv3d_1a_7x7(x)
        x = self.maxpool3d_2a_3x3(x)
        x = self.conv3d_2b_1x1(x)
        x = self.conv3d_2c_3x3(x)
        x = self.maxPool3d_3a_3x3(x)
        x = self.mixed_3b(x)
        x = self.mixed_3c(x)
        x = self.maxpool3d_4a_3x3(x)
        x = self.mixed_4b(x)
        x = self.mixed_4c(x)
        x = self.mixed_4d(x)
        x = self.mixed_4e(x)
        x = self.mixed_4f(x)
        x = self.maxpool3d_5a_2x2(x)
        x = self.mixed_5b(x)
        x = self.mixed_5c(x)
        # print(x.shape)
        x = self.avgpool3d(x, kernel_size=(2,7,7), strides=(1,1,1))
        return x
